// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

/// example:
/// bazel build //cc/appcoin2/support/feed:plot
/// ./bazel-bin/cc/appcoin2/support/feed/plot --driver ./data/coin2/feed/examples/...json

#include <boost/algorithm/string.hpp>
#include <cxxopts.hpp>

#include "coin2/app/driver.h"
#include "coin2/base/config/app_config.h"
#include "coin2/base/config/args.h"
#include "coin2/exchange/base/order/order.h"
#include "coin2/exchange/feed.pb.h"
#include "coin2/strategy/strategy.h"
#include "presto/quant/math/moving_average.h"
#define private public
#include "vitro.h"  // NOLINT

using coin2::exchange::base::feed::FeedUpdate;
using namespace coin2::app;
using namespace coin2::base::config;
using presto::math::TimeWindowMovingAverage;
using presto::math::SimpleMovingWindow;

struct BookStat {
  double bidp;
  double askp;
  std::map<int64_t, double> bps_qty;
  std::map<int64_t, double> bps_pq;
};

struct TradeStat {
  double trade_price;
  double trade_qty;
};

struct ProductData {
  explicit ProductData(int64_t window_size) {
    books_et.ResizeWindow(window_size);
    books.ResizeWindow(window_size);
    trades.ResizeWindow(window_size);
    trade_buys.ResizeWindow(window_size);
    trade_sells.ResizeWindow(window_size);
    markps.ResizeWindow(window_size);
    markps_et.ResizeWindow(window_size);
  }
  SimpleMovingWindow<BookStat> books_et;
  SimpleMovingWindow<BookStat> books;
  SimpleMovingWindow<TradeStat> trades;
  SimpleMovingWindow<TradeStat> trade_buys;
  SimpleMovingWindow<TradeStat> trade_sells;
  SimpleMovingWindow<double> markps;
  SimpleMovingWindow<double> markps_et;
};


struct Strategy : public ::coin2::strategy::IStrategy {
  std::optional<std::ofstream> outfile;
  std::string out_filename;
  int64_t prev_plot_time = 0L;
  int64_t plot_window_sec = 0L;
  int64_t last_timestamp = 0L;
  int64_t first_timestamp = 0L;
  std::map<std::string, ProductData> data;

  ProductData& GetItem(const FeedUpdate& upd) {
    auto id = upd.system_name() + ":" + upd.product().absolute_norm();
    if (data.count(id) == 0) {
      data.emplace(id, ProductData(plot_window_sec * 1e9));
    }
    return data.at(id);
  }

  void set_logfile(const std::string& outtxt_filename) {
    outfile = std::ofstream(outtxt_filename, std::ofstream::out);
  }

  void onMarkPriceFeed(const FeedUpdate& upd) override {
    auto& item = GetItem(upd);
    item.markps.Update(upd.timestamp(), upd.mark_price().mark_price);
    item.markps_et.Update(upd.mark_price().timestamp, upd.mark_price().mark_price);
  }

  void onTradeFeed(const FeedUpdate& upd) override {
    last_timestamp = upd.timestamp();
    if (!first_timestamp) first_timestamp = upd.timestamp();
    auto& item = GetItem(upd);

    TradeStat trade_stat;

    auto& trade = upd.trade();
    trade_stat.trade_price = trade.price;
    trade_stat.trade_qty = trade.fill_qty;

    item.trades.Update(upd.timestamp(), trade_stat);

    if (trade.sign() == 1) {
      item.trade_buys.Update(upd.timestamp(), trade_stat);
    } else if (trade.sign() == -1) {
      item.trade_sells.Update(upd.timestamp(), trade_stat);
    }
    (*outfile) << fmt::format(
          "T {} [{}] {} - [{}]{} @ {}\n",
          Iso8601FromTimestamp(upd.timestamp()),
          upd.timestamp(),
          upd.product().absolute_norm(),
          static_cast<int>(trade.sign()),
          trade.fill_qty,
          trade.price);
  }

  void onBookFeed(const FeedUpdate& upd) override {
    last_timestamp = upd.timestamp();
    if (!first_timestamp) first_timestamp = upd.timestamp();
    auto& book = upd.book();
    if (book.Ask0() && book.Bid0()) {
      BookStat book_stat;
      auto& item = GetItem(upd);

      auto ask = book.Ask0()->price;
      auto bid = book.Bid0()->price;
      book_stat.bidp = bid;
      book_stat.askp = ask;

      (*outfile) << fmt::format(
          "B {} [{}] {} - spd {:.2f}bps: {} - {}\n",
          Iso8601FromTimestamp(upd.timestamp()),
          upd.timestamp(),
          upd.product().absolute_norm(),
          ((ask - bid) / ask * 1e4),
          ask,
          bid);

      for (int range_bp : {10, 30}) {
        {
          double bound = book.Ask0()->price * (1 + range_bp * 1e-4);
          double sum = 0;
          double pq_sum = 0;
          for (auto& [p, q] : book.Asks()) {
            if (p > bound) {
              break;
            }
            sum += q;
            pq_sum += p * q;
          }
          book_stat.bps_qty[-range_bp] = sum;
          book_stat.bps_pq[-range_bp] = pq_sum;
        }
        {
          double bound = book.Bid0()->price * (1 - range_bp * 1e-4);
          double sum = 0;
          double pq_sum = 0;
          for (auto& [p, q] : book.Bids()) {
            if (p < bound) {
              break;
            }
            sum += q;
            pq_sum += p * q;
          }
          book_stat.bps_qty[range_bp] = sum;
          book_stat.bps_pq[range_bp] = pq_sum;
        }
      }
      item.books.Update(upd.timestamp(), book_stat);
      if (book.Timestamp() > 0) {
        item.books_et.Update(book.Timestamp(), book_stat);
      }

      if (upd.timestamp() - prev_plot_time > 10e9) {
        plot_book(upd.timestamp(), out_filename);
        prev_plot_time = upd.timestamp();
      }
    }
  }


  template <class MW, class Getter>
  auto get_pair(
      MW container,
      Getter getter) {
    std::vector<int64_t> tss;
    std::vector<double> vals;

    for (auto& elem : container) {
      tss.push_back(elem.first);
      vals.push_back(getter(elem.second));
    }
    return std::make_pair(tss, vals);
  }

  void plot_book(int64_t timestamp, const std::string& out_filename) {
    std::experimental::filesystem::path fpath(out_filename);
    Figure fig(3, 1);
    fig.title = fpath.stem();
    fig.width = 1000;
    fig.height = 800;
    fig.height_ratios = {5, 1, 1};

    {
      auto& ax = fig.axes(1, 1);
      ax.is_x_nanotimestamps = true;
      ax.grid = true;
      ax.legend = "lower right";
      std::string text;
      for (auto& [key, item] : data) {
        if (!item.markps.empty()) {
          auto [tss, vals] = get_pair(
              item.markps,
              [](const auto& elem) {
                  return elem;
              });
          auto& line = ax.line(":markp", tss, vals);
          line.alpha = 0.8;
          line.width = 0.4;
          line.color = "y";
          line.drawstyle = "steps-post";
        }
        if (!item.markps_et.empty()) {
          auto [ts_ets, vals] = get_pair(
              item.markps_et,
              [](const auto& elem) {
                  return elem;
              });
          auto& line = ax.line(":markp_et", ts_ets, vals);
          line.alpha = 0.8;
          line.width = 0.4;
          line.color = "k";
          line.drawstyle = "steps-post";
        }
        if (!item.books_et.empty()) {
          {
            auto [ts_ets, vals] = get_pair(
                item.books_et,
                [](const auto& elem) {
                    return elem.askp;
                });
            auto& line = ax.line(":askp_et", ts_ets, vals);
            line.alpha = 0.8;
            line.width = 0.4;
            line.color = "k";
            line.drawstyle = "steps-post";
          }
          {
            auto [ts_ets, vals] = get_pair(
                item.books_et,
                [](const auto& elem) {
                    return elem.bidp;
                });
            auto& line = ax.line(":bidp_et", ts_ets, vals);
            line.alpha = 0.8;
            line.width = 0.4;
            line.color = "k";
            line.drawstyle = "steps-post";
          }
        }
        if (!item.books.empty()) {
          {
            auto [tss, vals] = get_pair(
                item.books,
                [](const auto& elem) {
                    return elem.askp;
                });
            auto& line = ax.line(
                ":askp#" + fmt::format("{}", tss.size()),
                tss, vals);
            line.alpha = 0.8;
            line.width = 0.4;
            line.color = "r";
            line.drawstyle = "steps-post";
          }
          {
            auto [tss, vals] = get_pair(
                item.books,
                [](const auto& elem) {
                    return elem.bidp;
                });
            auto& line = ax.line(
                ":bidp#" + fmt::format("{}", tss.size()),
                tss, vals);
            line.alpha = 0.8;
            line.width = 0.4;
            line.color = "g";
            line.drawstyle = "steps-post";
          }
        }
        if (!item.trades.empty()) {
          {
            auto [tss, vals] = get_pair(
                item.trades,
                [](const auto& elem) {
                    return elem.trade_price;
                });
            auto& scatter = ax.scatter(
                ":trd_#" + fmt::format("{}", tss.size()),
                tss, vals);
            scatter.marker_size = 0.5;
            scatter.width = 0.5;
            scatter.alpha = 0.8;
            scatter.color = "g";
          }
          {
            auto [tss, vals] = get_pair(
                item.trade_sells,
                [](const auto& elem) {
                    return elem.trade_price;
                });
            auto& scatter = ax.scatter(
                ":trd_#" + fmt::format("{}", tss.size()),
                tss, vals);
            scatter.marker_size = 0.5;
            scatter.width = 0.5;
            scatter.alpha = 0.8;
            scatter.color = "r";
          }
        }
      }
    }
    {
      auto& ax = fig.axes(2, 1);
      ax.title = "Volume (Book Pq sum within 10bps)";
      ax.is_x_nanotimestamps = true;
      ax.grid = true;
      for (auto& [key, item] : data) {
        {
          auto [tss, vals] = get_pair(
              item.books,
              [](const auto& elem) {
                  return elem.bps_pq.at(-10);
              });
          auto& ask = ax.line(
              ":ask-10bps",
              tss,
              vals);
          ask.alpha = 0.8;
          ask.width = 0.4;
          ask.color = "r";
          ask.drawstyle = "steps-post";
        }
        {
          auto [tss, vals] = get_pair(
              item.books,
              [](const auto& elem) {
                  return elem.bps_pq.at(10);
              });
          auto& bid = ax.line(
              ":bid-10bps",
              tss,
              vals);
          bid.alpha = 0.8;
          bid.width = 0.4;
          bid.color = "g";
          bid.drawstyle = "steps-post";
        }
      }
    }
    {
      auto& ax = fig.axes(3, 1);
      ax.title = "Volume (Book Pq sum within 30bps)";
      ax.is_x_nanotimestamps = true;
      ax.grid = true;
      for (auto& [key, item] : data) {
        {
          auto [tss, vals] = get_pair(
              item.books,
              [](const auto& elem) {
                  return elem.bps_pq.at(-30);
              });
          auto& ask = ax.line(
              ":ask-30bps",
              tss,
              vals);
          ask.alpha = 0.8;
          ask.width = 0.4;
          ask.color = "r";
          ask.drawstyle = "steps-post";
        }
        {
          auto [tss, vals] = get_pair(
              item.books,
              [](const auto& elem) {
                  return elem.bps_pq.at(30);
              });
          auto& bid = ax.line(
              ":bid-30bps",
              tss,
              vals);
          bid.alpha = 0.8;
          bid.width = 0.4;
          bid.color = "g";
          bid.drawstyle = "steps-post";
        }
      }
    }
    Matplot matplot(fig);
    matplot.save(fpath);
    PyObject* pyplot = PyImport_ImportModule("matplotlib.pyplot");
    CHECK(pyplot);
    CHECK(PyObject_CallMethod(pyplot, "close", nullptr));
  }
};

using coin2::app::DriverConfig;


int main(int argc, char* argv[]) {
  /* ----------------------------------- opt ---------------------------------- */
  std::string out_filename;
  std::string outtxt_filename;

  DriverConfig app;

  cxxopts::Options opt("Feed Plot", "Plot feed");
  AddDriverOptions(&opt);                                         // --driver
  AddIntervalOptions(&opt);                                       // --start --end --duration
  opt.add_options()                                               // add options
      ("fastfeed", "use fastfeed")                                //
      ("live", "use live")
      ("raw", "use raw")                                          //
      ("plot_window_sec", "plot_window_sec", cxxopts::value<int64_t>()->default_value("6000000"))
      ("ws_host", "ws_host", cxxopts::value<std::string>()->default_value(""))
      ("out", "out", cxxopts::value<std::string>()->default_value("feedplot.png"))
      ("out_txt", "out_txt", cxxopts::value<std::string>()->default_value("pic/feed_out.txt"))
      ("machine", "feed machine", cxxopts::value<std::string>())  //
      ("mea", "M.E.A", cxxopts::value<std::string>())             //
      ("recipe", "recipe", cxxopts::value<std::string>())         //
      ("product", "a norm symbol", cxxopts::value<std::vector<std::string>>());  //

  auto res = opt.parse(argc, argv);

  if (!ParseDriverOptions(res, &app)) {
    auto& exchanges = *app.mutable_feed()->mutable_exchanges();
    auto& exchange = exchanges[res["mea"].as<std::string>()];

    exchange.set_mea(res["mea"].as<std::string>());
    exchange.set_recipe(res["recipe"].as<std::string>());
    for (const auto& norm : res["product"].as<std::vector<std::string>>()) {
      exchange.mutable_products()->add_norms(norm);
    }

    if (res["live"].as<bool>()) {
      exchange.set_type(coin2::exchange::feed::FeedSystemType::FEED_SYSTEM_TYPE_LIVE);
      exchange.mutable_websocket()->set_num_workers(1);
      if (!res["ws_host"].as<std::string>().empty()) {
        exchange.mutable_websocket()->set_ws_host(res["ws_host"].as<std::string>());
      }
    } else {
      exchange.set_type(coin2::exchange::feed::FeedSystemType::FEED_SYSTEM_TYPE_ARCHIVE);
      if (!(res["fastfeed"].as<bool>() ^ res["raw"].as<bool>())) {
        THROW() << "You should provide only one of --fastfeed or --raw";
      }
      bool use_fastfeed = res["fastfeed"].as<bool>();
      exchange.mutable_archive()->set_archive_type(
          use_fastfeed ? coin2::exchange::feed::ArchiveConfig::FASTFEED
                      : coin2::exchange::feed::ArchiveConfig::RAW);
      exchange.mutable_archive()->set_machine(res["machine"].as<std::string>());
    }
  }
  ParseIntervalOptions(
      res,
      app.mutable_feed()->mutable_common()->mutable_archive()->mutable_interval());

  out_filename = res["out"].as<std::string>();
  outtxt_filename = res["out_txt"].as<std::string>();

  /* ----------------------------------- run ---------------------------------- */
  printf("reading ..\n");

  Strategy strategy;
  strategy.set_logfile(outtxt_filename);
  strategy.out_filename = out_filename;
  strategy.plot_window_sec = res["plot_window_sec"].as<int64_t>();
  ::coin2::app::Driver driver(app, &strategy);
  driver.Run();
  strategy.plot_book(strategy.last_timestamp, out_filename);
  return 0;
}
